# -*- coding: utf-8 -*-
import sys, struct

from ..libmoops import TEXT_ADDRESS, DATA_ADDRESS, debug_print, verbose_print

class TranslationError( Exception ): pass

BIG_ENDIAN = False
USE_PSEUDO = True

STRUCT_FORMAT='>' if BIG_ENDIAN else '<'

opcodes = {
    'add':'000000' , 'addu':'000000' , 'and':'000000' , 'break':'000000' , 'div':'000000' , 'divu':'000000' ,
    'jalr':'000000' , 'jr':'000000', 'mfhi':'000000' , 'mflo':'000000' , 'mthi':'000000' , 'mtlo':'000000' ,
    'mult':'000000' , 'multu':'000000' , 'nor':'000000' , 'or':'000000' , 'sll':'000000' , 'sllv':'000000' ,
    'slt':'000000' , 'sltu':'000000' , 'sra':'000000' , 'srav':'000000' , 'srl':'000000' ,'srlv':'000000' ,
    'sub':'000000' , 'subu':'000000' , 'syscall':'000000' , 'xor':'000000' , 'rfe':'000000',

    'addi':'001000' , 'addiu':'001001' , 'andi':'001100' , 'beq':'000100' , 'bgez':'000001' , 'bgtz':'000111' ,
    'blez':'000110' , 'bltz':'000001', 'bne':'000101' , 'lb':'100000' , 'lbu':'100100' , 'lh':'100001' ,
    'lhu':'100101' , 'lui':'001111' , 'lw':'100011' , 'lwc1':'110001' , 'ori':'001101' , 'sb':'101000' ,
    'slti':'001010' , 'sltiu':'001011' , 'sh':'101001' , 'sw':'101011' , 'swc1':'111001' , 'xori':'001110' ,

    'j':'000010' , 'jal':'000011' ,

    'add.s':'010001' , 'cvt.s.w':'010001' , 'cvt.w.s':'010001' , 'div.s':'010001' , 'mfc0':'010000' , 'mfc1':'010001' , 'abs.s':'010001', 'abs.d':'010001', 'neg.s':'010001', 'neg.d':'010001',
    'mov.s':'010001' , 'mtc0':'010000', 'mtc1':'010001' , 'mul.s':'010001' , 'sub.s':'010001', 'c.eq.s':'010001' , 'c.eq.d':'010001' ,
    'c.lt.s':'010001' , 'c.lt.d':'010001' , 'c.le.s':'010001' , 'c.le.d':'010001' ,

    'bc1t':'010001' ,  'bc1f':'010001' , 'bc1tl':'010001' , 'bc1fl':'010001'

}

functs = {
    'add':'100000' , 'addu':'100001' , 'and':'100100' , 'break':'001101' , 'div':'011010' , 'divu':'011011' ,
    'jalr':'001001' , 'jr':'001000', 'mfhi':'010000' , 'mflo':'010010' , 'mthi':'010001' , 'mtlo':'010011' ,
    'mult':'011000' , 'multu':'011001' , 'nor':'100111' , 'or':'100101' , 'sll':'000000' ,
    'sllv':'000100' , 'slt':'101010' , 'sltu':'101011' , 'sra':'000011' , 'srav':'000111' , 'srl':'000010' ,
    'srlv':'000110' , 'sub':'100010' , 'subu':'100011' , 'syscall':'001100' , 'xor':'100110' , 'rfe':'010000' ,

    'add.s':'000000' , 'cvt.s.w':'100000' , 'cvt.w.s':'100100' , 'div.s':'000011' , 'mfc0':'000000' , 'mfc1':'000000' , 'abs.s':'000101' , 'abs.d':'000101' , 'neg.s':'000111' , 'neg.d':'000111' ,
    'mov.s':'000110' , 'mtc0':'0000000' , 'mtc1':'000000' , 'mul.s':'000010' , 'sub.s':'000001' , 'c.eq.s':'110010' ,
    'c.eq.d':'110010' , 'c.lt.s':'111100' , 'c.lt.d':'111100' , 'c.le.s':'111110' , 'c.le.d':'111110' ,
    'bc1t':'00001' ,  'bc1f':'00000' , 'bc1tl':'00011' , 'bc1fl':'00010'
}

forms = {
    'cvt.s.w':'10100' , 'cvt.w.s':'10000' , 'mfc0':'00000' , 'mfc1':'00000' , 'mtc0':'00100' , 'mtc1':'00100' ,
    'bc1t':'01000' ,  'bc1f':'01000' , 'bc1tl':'01000' , 'bc1fl':'01000'
}

reg_names=['zero', 'at', 'v0', 'v1', 'a0', 'a1', 'a2', 'a3', 't0', 't1', 't2', 't3', 't4', 't5', 't6', 't7',
             's0', 's1', 's2', 's3', 's4', 's5', 's6', 's7', 't8', 't9', 'k0', 'k1', 'gp', 'sp', 'fp', 'ra']

jump_table={}
data_store=bytes()


def format_instruction( asm_code ):
    asm_code = asm_code.split(None, 1)
    if len(asm_code)>1: asm_code=[asm_code[0]]+[x.strip() for x in asm_code[1].replace(',', ' ').split()]
    return asm_code

def splitbits( x ):
    return (int(x[:16], 2), int(x[16:], 2))

def float2int( x ):
    return str(struct.unpack(STRUCT_FORMAT+'I', struct.pack(STRUCT_FORMAT+'f', float(x)))[0])

def int2float( x ):
    return struct.unpack(STRUCT_FORMAT+'f', struct.pack(STRUCT_FORMAT+'I', int(x)))[0]

def depseudo( asm_code ):
    op=format_instruction( asm_code )
    if op[0]=='nop': return ("sll $zero,$zero,0",)

    if op[0] in ('div', 'divu', 'mul', 'mulu') and len(op)==4:
        if op[3][0]!='$':
            tmp=(depseudo("li $at,%s"%op[3]))
            op[3]='$at'
        else: tmp=()
        if op[0]=='mul': op[0]='mult'
        if op[0]=='mulu': op[0]='multu'
        return tmp+("%s %s,%s"%(op[0], op[2], op[3]), "mflo %s"%op[1])

    if op[0] in ('rem', 'remu'):
        if op[3][0]!='$':
            tmp=(depseudo("li $at,%s"%op[3]))
            op[3]='$at'
        else: tmp=()
        op[0]='divu' if op[0][-1]=='u' else 'div'
        return tmp+("%s %s,%s"%(op[0], op[2], op[3]), "mflo %s"%op[1])

    if op[0] in ('add', 'addu', 'sub', 'subu') and op[3][0]!='$':
        if op[0] in ('sub', 'subu'):
            if op[3][0]=='-': op[3]=op[3][1:]
            else: op[3]='-'+op[3]
            op[0]='add' if op[0]=='sub' else 'addu'
        if op[0][-1]=='u': op[0]=op[0][:-1]+'iu '
        else: op[0]+='i '

        return (op[0]+', '.join(op[1:]),)

    if op[0] in ('beq', 'bne') and op[2][0]!='$':
        hi, lo=splitbits(format_immediate(op[2], 32))
        if hi: return ("lui $at,%s"%hi, "ori $at,$zero,%s"%lo, "%s %s, $at, %s"%(op[0], op[1], op[3]))
        return ("ori $at,$zero,%s"%lo, "%s %s, $at, %s"%(op[0], op[1], op[3]))

    if op[0] in ('blt', 'bge'):
        set_="slt%s $at,%s,%s"%('i' if op[2][0]!='$' else '', op[1], op[2])
        branch="%s $at,$zero,%s"%('bne' if op[0]=='blt' else 'beq', op[3])
        return (set_, branch)

    if op[0] in ('bgt', 'ble'):
        set_="slt%s $at,%s,%s"%('i' if op[2][0]!='$' else '', op[2], op[1])
        branch="%s $at,$zero,%s"%('bne' if op[0]=='bgt' else 'beq', op[3])
        return (set_, branch)

    if op[0] in ('mov', 'move'): return ("addu %s,$zero,%s"%(op[1], op[2]),)

    if op[0]=='not': return ("nor %s,%s,$zero"%(op[1], op[2]),)

    if op[0]=='la':
        if op[2][0] not in '-0123456789': hi=lo=op[2]
        else: hi, lo=splitbits(format_immediate(op[2], 32))
        if hi: return ("lui $at,%s"%hi, "ori %s,$at,%s"%(op[1], lo))
        return ("ori %s,$zero,%s"%(op[1], lo),)

    if op[0]=='li.s':
        hi, lo=splitbits(format_immediate(float2int(op[2]), 32))
        if hi: return ("lui $at,%s"%hi, "ori $at,$at,%s"%lo, "mtc1 $at,%s"%op[1])
        return ("ori $at,$zero,%s"%(lo), "mtc1 $at,%s"%op[1])

    if op[0]=='li':
        hi, lo=splitbits(format_immediate(op[2], 32))
        if hi: return ("lui %s,%s"%(op[1], hi), "ori %s,%s,%s"%(op[1], op[1], lo))
        return ("ori %s,$zero,%s"%(op[1], lo),)

    if op[0] in ('lw', 'lwc1') and len(op)==3 and '(' not in op[2]:
        if op[2][0] not in '-0123456789': hi=lo=op[2]
        else: hi, lo=splitbits(format_immediate(op[2], 32))
        if hi: return ("lui $at,%s"%hi, "ori $at,$at,%s"%lo, "%s %s,0($at)"%(op[0], op[1]))
        return ("ori $at,$zero,%s"%lo, "%s %s,0($at)"%(op[0], op[1]))

    if op[0]=='mul':
        if len(op)==3:
            return ("mult %s, %s"%(op[1], op[2]),)
        return ("mult %s, %s"%(op[2], op[3]), "mflo %s"%op[1])

    if op[0]=='b': return ("bgez $zero,%s"%op[1],)

    return (asm_code,)

def calc_jumps( instructions ):
    real_instructions=[]
    address=[TEXT_ADDRESS, DATA_ADDRESS]
    mode=None
    for line in instructions:
        label, _, inst=line.rpartition(':')
        if label: jump_table[label.strip()]=address[mode]
        if inst:
            if '.globl' in inst: continue
            if mode is not None and inst[0]!='.':
                print("%x %s"%(address[mode], inst))
                address[mode]+=handle_data(inst.strip()) if mode else 4
            if '.data' in inst: mode=1
            elif '.text' in inst: mode=0
            real_instructions.append(inst)

    return real_instructions

def handle_data( inst ):
    address=0
    ret=bytes()
    if inst:
        type_, _, data=inst.partition(" ")
        if type_=='.word':
            for datum in data.split(','): address+=4
            data=int(data)
            ret += struct.pack(STRUCT_FORMAT+'I', data)
        elif type_ in ('.ascii', '.asciiz'):
            print("LOL", data, "ROFL")
            address+=len(data)-(1 if type_=='.asciiz' else 2)
            ret+=bytes(data[1:-1].encode('ascii'))
            if type_=='.asciiz': ret+=b'\x00'
    print(ret, "%%")
    global data_store
    data_store+=ret
    return address

def decomment( asm_code ):
    return asm_code.replace(';', '#').partition('#')[0].strip()

def pad( string, length ):
    if len(string)<length: string='0'*(length-len(string))+string
    return string[len(string)-length:]

def format_address( target, current=0 ): #current is 0 for jumps
    if target[0] not in '-1234567890':
      target= str(int(jump_table[target])-current) # this will be reconverted to binary
    return format_immediate(target, 18 if current else 28)[:-2] #low 2 bits don't matter, word alignment

def registers( register ): # Converts registers to binary
    #if n < 0: raise ValueError, "must be a positive integer"
    if register[0]!='$': raise TranslationError("Got literal, register expected")
    if register[1].lower()=='f': return fregisters( register )
    if register[1] not in '0123456789' and USE_PSEUDO: reg_num=reg_names.index(register[1:])
    else: reg_num=int(register[1:])
    return pad(bin(reg_num)[2:], 5)

def fregisters( fregister ): # Converts registers to binary
    #if n < 0: raise ValueError, "must be a positive integer"
    if fregister[1].lower()!='f': return registers( fregister )
    return pad(bin(int(fregister[2:]))[2:], 5)

def handle_literal( string ):
    neg=string[0]=='-'
    if neg: string=string[1:]
    if string[0]=='0' and len(string)>1:
        if string[1]=='x': ret= int(string, 16)
        elif string[1]=='b': ret= string
        else: ret=int(string, 8)
    else: ret= int(string)
    return bin((1<<32)-ret if neg else ret)[2:] # cut off the 0b prefix, and two's complement negatives

def format_immediate( target , length=16 ):
    if target[0] not in '-1234567890':
        target= str(jump_table[target])
    return pad(handle_literal(target), length)

def convert( asm_code, address=0xdeadbeef ):

    def convert_i(asm_code):
        if len(asm_code)==3:
            if asm_code[0] in ('bgez', 'bgtz', 'blez', 'bltz'):
                rs = registers(asm_code[1])
                immediate = format_address(asm_code[2], address)
                rt = ('00001' if asm_code[0] == 'bgez' else '00000')
            elif asm_code[0] == 'lui': # lui
                rs='00000'
                rt = registers( asm_code[1] )
                if asm_code[2][0] not in '-0123456789': immediate= str(jump_table[asm_code[2]]>>16)
                else: immediate=asm_code[2]
                immediate = format_immediate(immediate)
            else: # ex lw        $2,0($1)
                rt = registers( asm_code[1] )
                immediate = format_immediate( asm_code[2][:asm_code[2].index('(')] or '0' )
                rs = registers(asm_code[2][asm_code[2].index('(')+1:asm_code[2].index(')')])
        if len(asm_code) == 4: # I Type
            if asm_code[0] in ('beq', 'bne'):
                rs = registers( asm_code[1] )
                rt = registers( asm_code[2] )
                immediate=format_address(asm_code[3], address)
            else:
                rs = registers( asm_code[2] )
                rt = registers( asm_code[1] )
                immediate = format_immediate(asm_code[3])
        return opcode + rs + rt + immediate

    def convert_r(asm_code):
        rs = rt = rd = sa = '00000'
        function = functs[asm_code[0]]
        if asm_code[0]=='rfe': rs='10000'
        if len(asm_code)==2: #ex 000000 0000000000 00100 00000 010010    mflo $4 R Type
            if asm_code[0]=='jr': rs = registers( asm_code[1] )
            else: rd = registers( asm_code[1] )
        elif len(asm_code)==3: # R type
            if asm_code[0] == 'jalr':
                rd = registers( asm_code[1] )
                rs = registers( asm_code[2] )
            else:
                rs = registers( asm_code[1] )     #
                rt = registers( asm_code[2] )
        elif len(asm_code) == 4: # R Type
            assert asm_code[0] not in ('div','mult','divu','multu')
            rd = registers( asm_code[1] )
            if asm_code[0] in ('sll', 'sra', 'srl'):
                rt = registers( asm_code[2] )
                sa = format_immediate(asm_code[3], 5)
            elif asm_code[0] in ('sllv', 'srav', 'srlv'):
                rt = registers( asm_code[2] )
                rs = registers( asm_code[3] )
            else:
                rs = registers( asm_code[2] )
                rt = registers( asm_code[3] )
        return opcode + rs + rt + rd + sa + function

    def convert_j(asm_code):
        # cut off the last 2 bits as instructions are word aligned
        target = format_address(asm_code[1])
        return opcode + target

    def convert_fp (asm_code):
        if asm_code[0].endswith('.s'): format = '10000'
        elif asm_code[0].endswith('.d'): format = '10001'
        else: format = forms[asm_code[0]]
        if opcode not in ('bc1t', 'bc1f', 'bc1tl', 'bc1fl'):
            function = functs[asm_code[0]]
            ft = fs = fd = '00000'
            if len(asm_code) == 4:
                fd = fregisters(asm_code[1])
                fs = fregisters(asm_code[2])
                ft = fregisters(asm_code[3])
            elif len(asm_code) == 3:
                if asm_code[0] in ('c.eq.s', 'c.eq.d', 'c.lt.s', 'c.lt.d', 'c.le.s', 'c.le.d'):
                    fs = fregisters(asm_code[1])
                    ft = fregisters(asm_code[2])
                else:
                    if asm_code[0] in ('mtc0', 'mfc0', 'mtc1', 'mfc1'):
                        ft = fregisters(asm_code[1])
                    else:
                        fd = fregisters(asm_code[1])
                    fs = fregisters(asm_code[2])
            return opcode + format + ft + fs + fd + function
        else:
            return opcode + format + function + format_address(asm_code[1], address)

    asm_code= format_instruction( asm_code )
    opcode = opcodes[asm_code[0]]
    if opcode == '000000': #r-type
        return convert_r(asm_code)
    elif opcode[:5] == '00001':#j-type
        return convert_j(asm_code)
    elif opcode[:4] == '0100':
        return convert_fp(asm_code)
    else: #i-type
        return convert_i(asm_code)

def trans_file( infile, outfile=None, use_pseudo=True ):
    if outfile is None: outfile = infile.rsplit('.', 1)[0]+'.bin'

    global USE_PSEUDO
    USE_PSEUDO = use_pseudo

    verbose_print("Output file is: %s" % outfile)
    with open(infile) as f:
        instructions=f.read()

    binary = make_binary( instructions )

    with open(outfile, 'wb') as f:
        f.write(binary)

def make_binary( instructions ):

    instructions=[decomment(line) for line in instructions.split('\n') if decomment(line)]
    real_instructions=[]
    global data_store

    if USE_PSEUDO:
        for op in instructions:
            label, _, line=op.rpartition(':')
            #if op=='.data': break
            #if op[0]=='.':continue
            if label: real_instructions.append(label+':')
            if line: real_instructions.extend(depseudo(line))
    else: real_instructions=instructions
    text=bytes()

    real_instructions = calc_jumps(real_instructions)
    address=TEXT_ADDRESS
    for op in real_instructions:
        if op.strip()[0]=='.': continue
        val=int(convert(op, address), 2)
        hex_val="%08x"%val
        debug_print('0x%08x 0x%s\t\t%s' % (address, hex_val, op))
        address+=4
        text+=struct.pack(STRUCT_FORMAT+"I", val)

    debug_print("Jump table:")
    for k, v in list(jump_table.items()):
        debug_print("%s %x" % (k, v))

    binary = struct.pack(STRUCT_FORMAT+"I", len(text)>>2)
    binary+= b'\x00\x00\x00\x00' #no .ktext yet
    binary+= text
    binary+= data_store

    verbose_print("Raw binary size %i" % len(binary))
    if len(binary)%4:
        binary+= b'\x00'*(4-(len(binary)%4))
    verbose_print("Padded binary size %i" % len(binary))

    #cleanup
    jump_table.clear()
    data_store=bytes()

    debug_print("Raw binary is:")
    debug_print(repr(binary))

    return binary

if __name__=='__main__':
    if len(sys.argv)==1: print("Usage: %s file1 [file2] ... [fileN]" % sys.argv[0])
    else:
        for filename in sys.argv[1:]:
            trans_file(filename)

